import pandas as pd

from engine.strategy import Strategy
from engine.datafeed import CSVDatafeed
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

instrument = '162411.SZ'


#  三重滤网策略
class TripleFilterStrategy(Strategy):
    def __init__(self, feed):
        super(TripleFilterStrategy, self).__init__(feed)

        # 计算指标集 周线MACD， 日线KDJ等
        df = feed.get_df(instrument)
        df['close_weekly'] = df['close'].rolling(window=5).apply(lambda x:x[-1])
        df['high_weekly'] = df['high'].rolling(window=5).max()
        df['low_weekly'] = df['low'].rolling(window=5).min()

        #df.index = pd.to_datetime(df['date'])
        #se = df['close'].resample('W').last()
        print(df)

    def _calc_signal(self, df_bar):
        return 'BUY'  # SELL

    def onbar(self, index, date, df_bar):
        # 当买入信号时，全仓买入；卖出信号则平仓；

        signal = self._calc_signal(df_bar)
        if signal == 'BUY':
            self.acc.order_percent(instrument, 1.0)
        elif signal == 'SELL':
            self.acc.order_percent(instrument, 0.0)


feed = CSVDatafeed()
feed.add_data(instrument, '../datas/{}.csv'.format(instrument))

s = TripleFilterStrategy(feed=feed)
s.run()
s.plot(benchmark=[instrument])  # 需要在feed里加载地数据的
